import sys
sys.path.append('/home/llm_user/index/meta-learning/stable_meta_learning')
from stable_baselines3.common.vec_env import DummyVecEnv
import numpy as np
import csv

import torch

from stable_meta_learning.pearl import PEARL_SAC
from stable_meta_learning.envs import make_env,get_state,states_to_result


frictions = [1,2,3,4]
for friction in frictions:
    
    test_env_1 = make_env('PandaPush-v3',lateral_friction=friction)
    test_env = DummyVecEnv([test_env_1])
    model = PEARL_SAC.load('/home/llm_user/index/meta-learning/stable_meta_learning/checkpoints/PEARL_SAC1.zip',
                            env=test_env,
                            device=torch.device('cuda:0'),
    )
    out_csv = 'train1234-test10.csv'

    states = []
    success_num = 0

    pickandplace_num = 0
    roll_num = 0
    push_num = 0

    for _ in range(100):
        model.encoder.clear_z()
        observations = test_env.reset()
        actions = None
        rewards = None
        _states = []
        while True:
            _states.append(get_state(test_env))
            if actions is not None:
                obs = []
                for key in observations:
                    if key == 'task_z': continue
                    obs.append(observations[key])
                obs = np.concatenate(obs,axis=-1)
                model.update_context(obs, actions, rewards)
                model.encoder.infer_posterior(model.encoder.context)
                observations['task_z'] = model.encoder.task_z.cpu().detach().numpy()
            
            actions, _ = model.predict(
                observations,  # type: ignore[arg-type]
                deterministic=True,
            )

            observations, rewards, dones, infos = test_env.step(actions)
            if dones[0]:
                # print('->'.join(states),end=' ')
                if infos[0]['is_success']:
                    _states.append('success')
                    print('success')
                    success_num += 1
                else:
                    _states.append('fail')
                    print('fail')
                break
        states.append(_states)

    with open(out_csv, 'w', newline='',encoding='utf-8') as file:
        writer = csv.writer(file)
        for _states in states:
            writer.writerow(_states)
            result = states_to_result(_states)
            if result == 'push':
                push_num += 1
            elif result == 'roll':
                roll_num += 1
            else:
                pickandplace_num += 1


    print('success_rate:', success_num / 100)
    print('pickandplace_rate:',pickandplace_num / 100)
    print('roll_rate:',roll_num / 100)
    print('push_rate:',push_num / 100)

    test_env.close()